from functools import partial

import numpy as np

from bm3d import bm3d

class Preprocessor:


    key_words = ["name", "params"]


    def __init__(self, **kwargs):
        if len(kwargs) == 0:
            self.name = None
            self.forward = None
            return

        self.name = kwargs["name"]
        self.params = kwargs["params"]
        self.forward = None

        if self.name is None:
            pass
        
        if self.name == "BM3D":
            self.forward = partial(bm3d, sigma_psd=self.params["sigma_psd"])


    def process(self, img):
        if self.forward is None:
            return img

        if self.name == "BM3D":
            img = np.expand_dims(img, axis=2)
            img = self.forward(img)
            img = img.squeeze()
            return img
            
        else:
            return self.forward(img)

    